package chapter07;

import org.apache.hadoop.hive.ql.exec.UDAF;
import org.apache.hadoop.hive.ql.exec.UDAFEvaluator;
import org.apache.hadoop.io.IntWritable;
import org.apache.hadoop.io.Text;

import java.io.IOException;

/**
 * @author yuhui
 * descriptions: 对Hive表中某一列的平均值
 * @create 2024 - 12 - 12 3:05 下午
 */
public class AvgUDAF extends UDAF {

    public static class AvgUDAFEvaluator implements UDAFEvaluator {

        private boolean isEmpty;
        private double sum;
        private long count;

        public AvgUDAFEvaluator() {
            super();
            init();
        }

        public void init() {
            isEmpty = true;
            sum = 0;
            count = 0;
        }

        public boolean iterate(IntWritable value) throws IOException {
            if (value == null) {
                return true;
            }
            sum += value.get();
            count++;
            isEmpty = false;
            return true;
        }

        public Text terminatePartial() throws IOException {
            if (isEmpty) {
                return new Text("0");
            } else {
                return new Text(String.valueOf(sum));
            }
        }

        public boolean merge(Text partial) throws IOException {
            if (partial == null) {
                return true;
            }
            String[] parts = partial.toString().split(":");
            sum += Long.parseLong(parts[0]);
            count += Long.parseLong(parts[1]);
            return true;
        }

        public Text terminate() throws IOException {
            if (count == 0) {
                return new Text("0");
            } else {
                return new Text(String.valueOf(sum / count));
            }
        }
    }
}